package com.example.yolov8android;

import android.content.Context;
import android.graphics.Bitmap;
import org.tensorflow.lite.DataType;
import org.tensorflow.lite.Interpreter;
import org.tensorflow.lite.support.common.FileUtil;
import org.tensorflow.lite.support.common.ops.NormalizeOp;
import org.tensorflow.lite.support.image.ImageProcessor;
import org.tensorflow.lite.support.image.TensorImage;
import org.tensorflow.lite.support.tensorbuffer.TensorBuffer;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;

/**
 * YOLOv8 目标检测器
 * 用于加载 TFLite 模型并执行目标检测
 */
public class Detector {
    // 基本配置参数
    private static final float CONFIDENCE_THRESHOLD = 0.25F;  // 置信度阈值
    private static final float IOU_THRESHOLD = 0.4F;  // 非极大值抑制阈值

    // 核心组件
    private final Context context;
    private final String modelPath;
    private final String labelPath;
    private Interpreter interpreter;
    private final List<String> labels = new ArrayList<>();
    private final ImageProcessor imageProcessor;
    
    // 模型尺寸信息
    private int inputWidth = 0;
    private int inputHeight = 0;
    private int numClasses = 0;
    private int numBoxes = 0;

    /**
     * 构造函数
     */
    public Detector(Context context, String modelPath, String labelPath) {
        this.context = context;
        this.modelPath = modelPath;
        this.labelPath = labelPath;
        
        // 创建图像预处理器 - 归一化图像
        this.imageProcessor = new ImageProcessor.Builder()
                .add(new NormalizeOp(0f, 255f))
                .build();
    }

    /**
     * 初始化检测器
     */
    public void setup() {
        try {
            // 加载模型
            ByteBuffer modelBuffer = FileUtil.loadMappedFile(context, modelPath);
            Interpreter.Options options = new Interpreter.Options();
            options.setNumThreads(4);
            interpreter = new Interpreter(modelBuffer, options);

            // 获取模型输入输出尺寸
            int[] inputShape = interpreter.getInputTensor(0).shape();
            int[] outputShape = interpreter.getOutputTensor(0).shape();
            inputWidth = inputShape[1];
            inputHeight = inputShape[2];
            numClasses = outputShape[1];
            numBoxes = outputShape[2];

            // 加载标签
            loadLabels();
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    /**
     * 加载标签文件
     */
    private void loadLabels() {
        try {
            InputStream inputStream = context.getAssets().open(labelPath);
            BufferedReader reader = new BufferedReader(new InputStreamReader(inputStream));
            String line;
            while ((line = reader.readLine()) != null && !line.isEmpty()) {
                labels.add(line);
            }
            reader.close();
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    /**
     * 执行目标检测
     * @param image 输入图像
     * @return 检测到的边界框列表
     */
    public List<BoundingBox> detect(Bitmap image) {
        try {
            // 调整图像大小并预处理
            Bitmap resizedImage = Bitmap.createScaledBitmap(image, inputWidth, inputHeight, false);
            TensorImage tensorImage = new TensorImage(DataType.FLOAT32);
            tensorImage.load(resizedImage);
            TensorImage processedImage = imageProcessor.process(tensorImage);

            // 创建输出缓冲区并运行模型
            TensorBuffer output = TensorBuffer.createFixedSize(
                new int[]{1, numClasses, numBoxes}, DataType.FLOAT32);
            interpreter.run(processedImage.getBuffer(), output.getBuffer());

            // 处理检测结果
            return processDetections(output.getFloatArray());
        } catch (Exception e) {
            e.printStackTrace();
            return new ArrayList<>();
        }
    }

    /**
     * 处理检测结果
     */
    private List<BoundingBox> processDetections(float[] detections) {
        List<BoundingBox> boxes = new ArrayList<>();

        // 处理每个检测框
        for (int i = 0; i < numBoxes; i++) {
            // 寻找最高置信度的类别
            float maxConf = -1.0f;
            int maxClassId = -1;
            
            for (int classId = 0; classId < numClasses - 4; classId++) {
                float conf = detections[i + numBoxes * (classId + 4)];
                if (conf > maxConf) {
                    maxConf = conf;
                    maxClassId = classId;
                }
            }

            // 如果置信度超过阈值，创建边界框
            if (maxConf > CONFIDENCE_THRESHOLD) {
                String className = labels.get(maxClassId);
                
                // 获取边界框坐标 (中心点坐标+宽高)
                float cx = detections[i];
                float cy = detections[i + numBoxes];
                float w = detections[i + numBoxes * 2];
                float h = detections[i + numBoxes * 3];
                
                // 转换为左上坐标和宽高
                float x1 = cx - (w/2F);
                float y1 = cy - (h/2F);
                
                // 边界检查 (值范围0-1)
                if (x1 >= 0 && x1 <= 1 && y1 >= 0 && y1 <= 1 && 
                    w >= 0 && w <= 1 && h >= 0 && h <= 1) {
                    // 创建实际的像素坐标的边界框对象(整数类型)
                    int actualX = Math.round(x1 * inputWidth);
                    int actualY = Math.round(y1 * inputHeight);
                    int actualWidth = Math.round(w * inputWidth);
                    int actualHeight = Math.round(h * inputHeight);
                    
                    boxes.add(new BoundingBox(
                        actualX, actualY, actualWidth, actualHeight,
                        maxConf, maxClassId, className
                    ));
                }
            }
        }

        // 应用非极大值抑制，消除重叠框
        return applyNMS(boxes);
    }

    /**
     * 应用非极大值抑制(NMS)
     */
    private List<BoundingBox> applyNMS(List<BoundingBox> boxes) {
        if (boxes.isEmpty()) {
            return new ArrayList<>();
        }
        
        // 按置信度排序
        List<BoundingBox> sortedBoxes = new ArrayList<>(boxes);
        Collections.sort(sortedBoxes, (b1, b2) -> Float.compare(b2.getCnf(), b1.getCnf()));
        
        List<BoundingBox> selectedBoxes = new ArrayList<>();
        List<BoundingBox> remainingBoxes = new ArrayList<>(sortedBoxes);

        // 选择置信度最高的框，移除重叠框
        while (!remainingBoxes.isEmpty()) {
            BoundingBox bestBox = remainingBoxes.remove(0);
            selectedBoxes.add(bestBox);
            
            remainingBoxes.removeIf(box -> calculateIoU(bestBox, box) >= IOU_THRESHOLD);
        }

        return selectedBoxes;
    }

    /**
     * 计算两个边界框的交并比(IoU)
     */
    private float calculateIoU(BoundingBox box1, BoundingBox box2) {
        // 获取两个框的坐标
        int x1 = Math.max(box1.getX(), box2.getX());
        int y1 = Math.max(box1.getY(), box2.getY());
        int x2 = Math.min(box1.getX() + box1.getWidth(), box2.getX() + box2.getWidth());
        int y2 = Math.min(box1.getY() + box1.getHeight(), box2.getY() + box2.getHeight());
        
        // 计算交集面积
        int intersectionWidth = Math.max(0, x2 - x1);
        int intersectionHeight = Math.max(0, y2 - y1);
        int intersectionArea = intersectionWidth * intersectionHeight;
        
        // 计算两个框的面积
        int box1Area = box1.getWidth() * box1.getHeight();
        int box2Area = box2.getWidth() * box2.getHeight();
        
        // 计算 IoU = 交集面积 / (两个框面积之和 - 交集面积)
        return intersectionArea / (float)(box1Area + box2Area - intersectionArea);
    }

    /**
     * 获取模型输入宽度
     */
    public int getInputWidth() {
        return inputWidth;
    }

    /**
     * 获取模型输入高度
     */
    public int getInputHeight() {
        return inputHeight;
    }
} 